package com.jianying.day08;

import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;

/**
 * ClassName: Flink01_KeyedState_ValueState_1
 * Package: com.atguigu.day08
 * Description:
 *
 * @Author: tubage
 * @Create: 2024/4/9 09:40
 * @Version: 1.0
 */
public class Flink08_OpetatorState_ListState {
    public static void main(String[] args) throws Exception {
        // TODO 在map算子每个并行度中计算 数据个数
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);
        DataStreamSource<String> socketDS = env.socketTextStream("localhost", 8888);

        socketDS.process(
                new MyProcess()
                // new ProcessFunction<String, String>() {
                //     // TODO 普通变量也可以，作用范围也是每个并行度 但不能持久化
                //     Integer count = 0;
                //     @Override
                //     public void processElement(String value, ProcessFunction<String, String>.Context ctx, Collector<String> out) throws Exception {
                //         count++;
                //         out.collect("当前并行度："+ getRuntimeContext().getIndexOfThisSubtask() + ",数据个数：" + count);
                //
                //     }
                // }

        ).print();


        env.execute();
    }
}

class MyProcess extends ProcessFunction<String,String> implements CheckpointedFunction {
    Integer count = 0;
    ListState<Integer> countListState;
    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        System.out.println("snapshotState~~~");
        // TODO 清空
        countListState.clear();
        // TODO 要备份的数据放到状态中
        countListState.add(count);

    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        System.out.println("initializeState~~~");
        // TODO 状态初始化 初始化本地变量：程序启动和恢复时， 从状态中 把数据添加到 本地变量，每个子任务调用一次
        ListStateDescriptor<Integer> listStateDescriptor = new ListStateDescriptor<>("countListState", Integer.class);
        countListState = context.getOperatorStateStore().getListState(listStateDescriptor);
        if(context.isRestored()) {
            count = countListState.get().iterator().next();
        }
    }

    @Override
    public void processElement(String value, ProcessFunction<String, String>.Context ctx, Collector<String> out) throws Exception {
        count++;
        out.collect("当前并行度："+ getRuntimeContext().getIndexOfThisSubtask() + ",数据个数：" + count);
    }
}